{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SageMaker V3 Hyperparameter Training Example\n",
    "\n",
    "This notebook demonstrates hyperparameter handling in SageMaker V3 ModelTrainer using JSON and YAML files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import yaml\n",
    "import tempfile\n",
    "import shutil\n",
    "\n",
    "from sagemaker.train.model_trainer import ModelTrainer\n",
    "from sagemaker.train.configs import SourceCode\n",
    "from sagemaker.core.helper.session_helper import Session, get_execution_role"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Setup Session and Create Test Files\n",
    "\n",
    "Initialize the SageMaker session and create the hyperparameter files and training script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker_session = Session()\n",
    "role = get_execution_role()\n",
    "\n",
    "# Expected hyperparameters\n",
    "EXPECTED_HYPERPARAMETERS = {\n",
    "    \"integer\": 1,\n",
    "    \"boolean\": True,\n",
    "    \"float\": 3.14,\n",
    "    \"string\": \"Hello World\",\n",
    "    \"list\": [1, 2, 3],\n",
    "    \"dict\": {\n",
    "        \"string\": \"value\",\n",
    "        \"integer\": 3,\n",
    "        \"float\": 3.14,\n",
    "        \"list\": [1, 2, 3],\n",
    "        \"dict\": {\"key\": \"value\"},\n",
    "        \"boolean\": True,\n",
    "    },\n",
    "}\n",
    "\n",
    "DEFAULT_CPU_IMAGE = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310\"\n",
    "\n",
    "# Create temporary directory\n",
    "temp_dir = tempfile.mkdtemp()\n",
    "source_dir = os.path.join(temp_dir, \"source\")\n",
    "os.makedirs(source_dir, exist_ok=True)\n",
    "\n",
    "print(f\"Created temporary directory: {temp_dir}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Create Hyperparameter Files and Training Script\n",
    "\n",
    "Create JSON and YAML hyperparameter files and a training script that validates them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create JSON hyperparameters file\n",
    "json_file = os.path.join(source_dir, \"hyperparameters.json\")\n",
    "with open(json_file, 'w') as f:\n",
    "    json.dump(EXPECTED_HYPERPARAMETERS, f, indent=2)\n",
    "\n",
    "# Create YAML hyperparameters file\n",
    "yaml_file = os.path.join(source_dir, \"hyperparameters.yaml\")\n",
    "with open(yaml_file, 'w') as f:\n",
    "    yaml.dump(EXPECTED_HYPERPARAMETERS, f, default_flow_style=False, indent=2)\n",
    "\n",
    "print(\"Created hyperparameter files\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create training script that validates hyperparameters\n",
    "training_script = '''\n",
    "import argparse\n",
    "import json\n",
    "import os\n",
    "\n",
    "EXPECTED_HYPERPARAMETERS = {\n",
    "    \"integer\": 1,\n",
    "    \"boolean\": True,\n",
    "    \"float\": 3.14,\n",
    "    \"string\": \"Hello World\",\n",
    "    \"list\": [1, 2, 3],\n",
    "    \"dict\": {\n",
    "        \"string\": \"value\",\n",
    "        \"integer\": 3,\n",
    "        \"float\": 3.14,\n",
    "        \"list\": [1, 2, 3],\n",
    "        \"dict\": {\"key\": \"value\"},\n",
    "        \"boolean\": True,\n",
    "    },\n",
    "}\n",
    "\n",
    "def parse_args():\n",
    "    parser = argparse.ArgumentParser(description=\"Test Hyperparameters\")\n",
    "    parser.add_argument(\"--string\", type=str, required=True)\n",
    "    parser.add_argument(\"--integer\", type=int, required=True)\n",
    "    parser.add_argument(\"--float\", type=float, required=True)\n",
    "    parser.add_argument(\"--boolean\", type=lambda x: json.loads(x), required=True)\n",
    "    parser.add_argument(\"--list\", type=lambda x: json.loads(x), required=True)\n",
    "    parser.add_argument(\"--dict\", type=lambda x: json.loads(x), required=True)\n",
    "    return parser.parse_args()\n",
    "\n",
    "def main():\n",
    "    args = parse_args()\n",
    "    print(f\"Received hyperparameters: {args}\")\n",
    "\n",
    "    # Validate hyperparameters\n",
    "    assert args.string == EXPECTED_HYPERPARAMETERS[\"string\"]\n",
    "    assert args.integer == EXPECTED_HYPERPARAMETERS[\"integer\"]\n",
    "    assert args.boolean == EXPECTED_HYPERPARAMETERS[\"boolean\"]\n",
    "    assert args.float == EXPECTED_HYPERPARAMETERS[\"float\"]\n",
    "    assert args.list == EXPECTED_HYPERPARAMETERS[\"list\"]\n",
    "    assert args.dict == EXPECTED_HYPERPARAMETERS[\"dict\"]\n",
    "\n",
    "    # Validate environment variables\n",
    "    params = json.loads(os.environ[\"SM_HPS\"])\n",
    "    print(f\"SM_HPS: {params}\")\n",
    "    assert params == EXPECTED_HYPERPARAMETERS\n",
    "\n",
    "    print(\"All hyperparameter validations passed!\")\n",
    "    \n",
    "    # Save results\n",
    "    model_dir = os.environ.get(\"SM_MODEL_DIR\", \"/opt/ml/model\")\n",
    "    os.makedirs(model_dir, exist_ok=True)\n",
    "    \n",
    "    results = {\"status\": \"success\", \"hyperparameters\": params}\n",
    "    with open(os.path.join(model_dir, \"results.json\"), \"w\") as f:\n",
    "        json.dump(results, f, indent=2)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n",
    "'''\n",
    "\n",
    "with open(os.path.join(source_dir, \"train.py\"), 'w') as f:\n",
    "    f.write(training_script)\n",
    "\n",
    "# Create requirements file\n",
    "with open(os.path.join(source_dir, \"requirements.txt\"), 'w') as f:\n",
    "    f.write(\"omegaconf\\n\")\n",
    "\n",
    "print(\"Created training script and requirements\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Training with JSON Hyperparameters\n",
    "\n",
    "Train a model using hyperparameters loaded from a JSON file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_code = SourceCode(\n",
    "    source_dir=source_dir,\n",
    "    requirements=\"requirements.txt\",\n",
    "    entry_script=\"train.py\",\n",
    ")\n",
    "\n",
    "json_trainer = ModelTrainer(\n",
    "    sagemaker_session=sagemaker_session,\n",
    "    training_image=DEFAULT_CPU_IMAGE,\n",
    "    hyperparameters=json_file,\n",
    "    source_code=source_code,\n",
    "    base_job_name=\"hp-contract-hyperparameter-json\",\n",
    ")\n",
    "\n",
    "print(\"ModelTrainer created with JSON hyperparameters!\")\n",
    "print(f\"Hyperparameters loaded: {json_trainer.hyperparameters}\")\n",
    "\n",
    "# Verify hyperparameters match expected values\n",
    "assert json_trainer.hyperparameters == EXPECTED_HYPERPARAMETERS\n",
    "print(\"✓ JSON hyperparameters match expected values!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Starting training with JSON hyperparameters...\")\n",
    "\n",
    "json_trainer.train()\n",
    "print(f\"JSON hyperparameters training completed: {json_trainer._latest_training_job.training_job_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Training with YAML Hyperparameters\n",
    "\n",
    "Train a model using hyperparameters loaded from a YAML file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yaml_trainer = ModelTrainer(\n",
    "    sagemaker_session=sagemaker_session,\n",
    "    training_image=DEFAULT_CPU_IMAGE,\n",
    "    hyperparameters=yaml_file,\n",
    "    source_code=source_code,\n",
    "    base_job_name=\"hp-contract-hyperparameter-yaml\",\n",
    ")\n",
    "\n",
    "print(\"ModelTrainer created with YAML hyperparameters!\")\n",
    "print(f\"Hyperparameters loaded: {yaml_trainer.hyperparameters}\")\n",
    "\n",
    "# Verify hyperparameters match expected values\n",
    "assert yaml_trainer.hyperparameters == EXPECTED_HYPERPARAMETERS\n",
    "print(\"✓ YAML hyperparameters match expected values!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Starting training with YAML hyperparameters...\")\n",
    "\n",
    "yaml_trainer.train()\n",
    "print(f\"YAML hyperparameters training completed: {yaml_trainer._latest_training_job.training_job_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Compare Training Results\n",
    "\n",
    "Compare the results from both hyperparameter approaches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_jobs = [\n",
    "    (\"JSON File\", json_trainer),\n",
    "    (\"YAML File\", yaml_trainer)\n",
    "]\n",
    "\n",
    "print(\"Training Job Comparison:\")\n",
    "print(\"=\" * 40)\n",
    "\n",
    "for approach_name, trainer in training_jobs:\n",
    "    job_name = trainer._latest_training_job.training_job_name\n",
    "    model_artifacts = trainer._latest_training_job.model_artifacts\n",
    "    \n",
    "    print(f\"\\n{approach_name}:\")\n",
    "    print(f\"  Job Name: {job_name}\")\n",
    "    print(f\"  Model Artifacts: {model_artifacts}\")\n",
    "    print(f\"  Status: Completed\")\n",
    "    \n",
    "    # Verify all hyperparameters are identical\n",
    "    assert trainer.hyperparameters == EXPECTED_HYPERPARAMETERS\n",
    "\n",
    "print(\"\\n✓ All training jobs completed successfully with identical hyperparameters!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: Clean Up\n",
    "\n",
    "Clean up temporary files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    shutil.rmtree(temp_dir)\n",
    "    print(f\"Cleaned up temporary directory: {temp_dir}\")\n",
    "except Exception as e:\n",
    "    print(f\"Could not clean up temp directory: {e}\")\n",
    "\n",
    "print(\"Cleanup completed!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This notebook demonstrated:\n",
    "1. **JSON hyperparameters**: Loading from JSON files\n",
    "2. **YAML hyperparameters**: Loading from YAML files\n",
    "3. **Validation**: Ensuring loaded hyperparameters match expected values\n",
    "4. **File-based configuration**: Managing hyperparameters as external files\n",
    "\n",
    "File-based hyperparameters provide better version control, reproducibility, and support for complex nested structures."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv-test",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
